#include "Socket.h"
#include <winsock2.h>
#include <boost/optional.hpp>
#include <boost/shared_ptr.hpp>
#include "utils.h"

namespace Broady {

bool Socket::s_init = false;

SocketPtr Socket::Create( void )
{
	if( !s_init )
	{
		struct WSAData wsdata;

		int error = WSAStartup( 0x0101, &wsdata );

		if( error != 0 )
		{
			OutputMessage( "Failed to initialize WinSock: %lu\n", GetLastError( ) );

			return SocketPtr( );
		}

		s_init = true;
	}

	SocketPtr s = SocketPtr( new Socket( ) );

	if( !s || !s->_Obtain( ) )
	{
		return SocketPtr( );
	}

	return s;
}

int Socket::LastError( void ) const
{
	return m_lastError;
}

SOCKET Socket::Handle( void ) const
{
	return m_handle;
}

bool Socket::IsValid( void ) const
{
	return m_handle != INVALID_SOCKET;
}

Socket::Socket( void )
	: m_lastError( ERROR_SUCCESS )
	, m_handle( INVALID_SOCKET )
{
}

Socket::~Socket( void )
{
	_Close( );
}

bool Socket::_Close( void )
{
	if( m_handle != INVALID_SOCKET )
	{
		if( shutdown( m_handle, SD_BOTH ) == SOCKET_ERROR )
		{
			_PrintError( "shutdown", m_lastError = WSAGetLastError( ) );
		}

		if( closesocket( m_handle ) == SOCKET_ERROR )
		{
			_PrintError( "closesocket", m_lastError = WSAGetLastError( ) );

			return false;
		}

		m_handle = INVALID_SOCKET;
	}

	return true;
}

bool Socket::_Obtain( void )
{
	SOCKET s = socket( PF_INET, SOCK_DGRAM, IPPROTO_UDP );

	if( s == INVALID_SOCKET )
	{
		_PrintError( "socket", m_lastError = WSAGetLastError( ) );

		return false;
	}

	_Close( );
	m_handle = s;

	return true;
}

bool Socket::SendTo( const void* data, unsigned int& len, unsigned long ip, unsigned short port )
{
	boost::optional<unsigned int> result = _sendto( m_handle, data, len, ip, port );

	if( !result )
	{
		return false;
	}

	if( len != *result )
	{
		OutputMessage( "Warning: SendTo() did not send everything. (%u/%u)\n", len, *result );
	}

	len = *result;

	return true;
}

boost::optional<unsigned int> Socket::Available( void )
{
	unsigned long arg = 1;

	if( ioctlsocket( m_handle, FIONREAD, &arg ) == SOCKET_ERROR )
	{
		_PrintError( "ioctlsocket", m_lastError = WSAGetLastError( ) );

		return boost::optional<unsigned int>( );
	}

	return static_cast<unsigned int>( arg );
}

bool Socket::SetBroadcast( bool broadcast )
{
	char c = broadcast ? 1 : 0;

	if( setsockopt( m_handle, SOL_SOCKET, SO_BROADCAST, &c, sizeof( c ) ) == SOCKET_ERROR )
	{
		_PrintError( "setsockopt", m_lastError = WSAGetLastError( ) );

		return false;
	}

	return true;
}

void Socket::_PrintError( const char* const name, int error )
{
	OutputMessage( "WinSock: %s() returned an error: %d\n", name, error );
}

boost::optional<unsigned int> Socket::_sendto( SOCKET s, const void* data, unsigned int len, unsigned long ip, unsigned short port )
{
	if( !data || len == 0 )
	{
		OutputMessage( "Error: No data passed to _sendto().\n" );

		return boost::optional<unsigned int>( );
	}

	struct sockaddr_in addr;

	addr.sin_family = AF_INET;
	addr.sin_port = htons( port );
	addr.sin_addr.s_addr = ip;

	int result = sendto( s, static_cast<const char*>( data ), len, 0, reinterpret_cast<const struct sockaddr*>( &addr ), sizeof( addr ) );

	if( result == SOCKET_ERROR )
	{
		_PrintError( "sendto", m_lastError = WSAGetLastError( ) );

		return boost::optional<unsigned int>( );
	}

	return static_cast<unsigned int>( result );
}

boost::optional<bool> Socket::RecvFrom( void* buffer, unsigned int& len, unsigned long* ip, unsigned short* port )
{
	boost::optional< std::pair<unsigned int, bool> > result = _recvfrom( m_handle, buffer, len, ip, port );

	if( !result )
	{
		return boost::optional<bool>( );
	}

	len = ( *result ).first;

	return ( *result ).second;
}

bool Socket::GetLocalAddress( unsigned long* ip, unsigned short* port )
{
	if( !ip && !port )
	{
		// Looks like we don't want anything...
		OutputMessage( "Error: No parameters at all for Socket::GetLocalAddress().\n" );

		return false;
	}

	struct sockaddr_in name;
	int namelen = sizeof( name );

	if( getsockname( m_handle, reinterpret_cast<struct sockaddr*>( &name ), &namelen ) == SOCKET_ERROR )
	{
		_PrintError( "getsockname", m_lastError = WSAGetLastError( ) );

		return false;
	}

	memcpy( ip, &name.sin_addr, sizeof( *ip ) );
	*port = htons( name.sin_port );

	return true;
}

bool Socket::Bind( unsigned long ip, unsigned short port )
{
	struct sockaddr_in addr;

	addr.sin_family = AF_INET;
	addr.sin_port = htons( port );
	addr.sin_addr.s_addr = ip;

	if( bind( m_handle, reinterpret_cast<const struct sockaddr*>( &addr ), sizeof( addr ) ) == SOCKET_ERROR )
	{
		_PrintError( "bind", m_lastError = WSAGetLastError( ) );

		return false;
	}

	return true;
}

// If there is no value from the boost::optional<T>, there was an error or no data available.
// Else, if the bool from std::pair<X,Y> is false, the buffer was not big enough, the unsigned int is the needed size.
// If the bool from std::pair<X,Y> is true, the unsigned int is the number of bytes read.
boost::optional< std::pair<unsigned int, bool> > Socket::_recvfrom(
		SOCKET s,
		void* buffer,
		unsigned int len,
		unsigned long* ip,
		unsigned short* port )
{
	boost::optional<unsigned int> avail = Available( );

	if( !avail )
	{
		return boost::optional< std::pair<unsigned int, bool> >( );
	}

	if( *avail == 0 )
	{
		// No data, no error
		return std::pair<unsigned int, bool>( 0, true );
	}

	if( *avail > len )
	{
		// UDP packets can only be read as a whole, buffer not large enough.
		return std::pair<unsigned int, bool>( *avail, false );
	}

	// Only check if we have a buffer after we check if we have data available.
	if( !buffer )
	{
		OutputMessage( "Error: No buffer passed to _recvfrom().\n" );

		return boost::optional< std::pair<unsigned int, bool> >( );
	}

	if( len > *avail )
	{
		// Buffer is large enough, but we don't need to read more than what's available.
		len = *avail;
	}

	struct sockaddr_in addr;
	int size = sizeof( addr );

	// Prepare the address to know who is sending the data.
	memset( &addr, 0, sizeof( addr ) );
	addr.sin_family = AF_INET;

	int r = recvfrom( s, static_cast<char*>( buffer ), len, 0, reinterpret_cast<struct sockaddr*>( &addr ), &size );

	if( r == SOCKET_ERROR )
	{
		_PrintError( "recvfrom", m_lastError = WSAGetLastError( ) );

		return boost::optional< std::pair<unsigned int, bool> >( );
	}

	if( ip )
	{
		*ip = addr.sin_addr.s_addr;
	}

	if( port )
	{
		*port = htons( addr.sin_port );
	}

	return std::pair<unsigned int, bool>( len, true );
}

} // namespace
